import os
import sys

# 获取当前脚本文件的绝对路径
current_file_path = os.path.abspath(__file__)
# 获取当前脚本文件的上级目录
parent_directory = os.path.dirname(os.path.dirname(current_file_path))
# 将上级目录添加到 sys.path
if parent_directory not in sys.path:
    sys.path.append(parent_directory)
import re
import sys
import aiofiles
import asyncio
import aiohttp
import time
import ssl
from config import config
from colorama import init, Fore
from openpyxl import Workbook
from yarl import URL
from aiohttp.resolver import AsyncResolver
from urllib.parse import urlencode, urlparse, urlunparse, parse_qs
sys.setrecursionlimit(10000)  # 将递归深度增加到 10000
plat=sys.platform
# 将http、https协议互相替换
def replace_protocol(url):
    if url.startswith('http:'):
        return url.replace('http:', 'https:')
    elif url.startswith('https:'):
        return url.replace('https:', 'http:')

# 异步文件生成器读取
async def read_lines(file_path, batch_size):
    def add_to_seen_urls(url):
        nonlocal count
        if count < batch_size and url not in seen_urls:
            urls.append(url)
            seen_urls.add(url)
            count += 1
    try:
        if batch_size < 20:
            batch_size = 10 #最小就是10，太小没必要了
        async with aiofiles.open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
            seen_urls = set() # 使用集合来存储已经出现过的 URL
            while True:
                urls = []
                count = 0  # 计数器，用于确保每批次不超过 batch_size
                while count < batch_size:
                    line = await file.readline() #在同一个 async with 上下文中每次调用 readline() 或 read() 等读取方法时文件指针会自动更新到下一行或下一个字节的位置。
                    if not line:
                        # 如果到达文件末尾，跳出循环
                        break
                    url = line.strip()
                    if url and not any(url.split('.')[-1].startswith(exclude_item) for exclude_item in config.exclude_e):

                        if config.replace_protocol:
                            if url.startswith("https://") or url.startswith("http://"):
                                url_replace_new=replace_protocol(url)
                                add_to_seen_urls(url_replace_new)
                                if count >= batch_size:
                                    break
                        if not url.startswith("https://") and not url.startswith("http://"):
                            https_url = "https://" + url
                            http_url = "http://" + url
                            # 添加 https 版本的 URL，如果它还没有被添加过
                            add_to_seen_urls(https_url)
                            add_to_seen_urls(http_url)
                            if count >= batch_size: # 如果达到 batch_size，则退出当前批次读取
                                break
                        else:
                            # 如果 URL 已有协议，则直接处理它
                            add_to_seen_urls(url)
                            if count >= batch_size: # 如果达到 batch_size，则退出当前批次读取
                                break
                # 如果有读取到行，则产生这批行
                if urls:
                    yield urls
                else:
                    break
    except Exception as e:
        print(e)
        print('文件路径或类型错误, 当前仅支持传递 txt 文本')
# 从字节数据块儿中解析title的标签
def gettitle(data):
    if data is None:
      title="未获取到标题"
      return title
    try:
        title = data.decode('utf-8').strip()  # 假设编码为UTF-8
    except:
        try:
            title = data.decode('gbk',errors='ignore').strip()
        except:
            title="未获取到标题"
    return title


# 单个url探测的函数
async def get_title_code(url,semaphore,method,session,redirect=True,retry=0):
    header = config.headers
    parsed_url = urlparse(url) #解析url
    # print("正在探测:"+url)
    # 解析URL，提取域名
    if config.autoUpdate_Host:
        host = parsed_url.hostname #通过urlparse解析后获取host
        header['Host']=host  # 根据请求的url自动设置 Host 头部
    # print(header)
    location = ''
    try:
        # 实现从config.py动态添加get或者post请求参数的功能
        existing_params = parse_qs(parsed_url.query)#解析 URL的现有查询参数为字典类型
        # 构造动态session的参数
        request_params = {
            "url": url,
            "headers": header,
            "ssl": False,
            "proxy": None,
            "allow_redirects": redirect,

        }
        # 添加GET或POST参数
        if method.lower() == 'get' and config.get_params is not None:
            existing_params.update(config.get_params) # 合并现有参数和新参数
            query_string = urlencode(existing_params, doseq=True)  # 将字典的get参数转化为url编码：允许将字典中相同的键对应多个值
            full_url = urlunparse(parsed_url._replace(query=query_string))# 使用 urlunparse 构造完整的 URL 替换原始url的参数部分
            url=full_url
            # 以上都是为了输出保存的适合可以直观看到get请求的参数，以下才是涉及到请求的操作
            request_params['params'] = config.get_params  # 仅在GET请求时添加
        elif method.lower() == 'post' and config.post_data is not None:
            # 自动判断是否为 JSON 数据
            if isinstance(config.post_data, dict):  # 检查是否为字典，是则以JSON格式发送
                request_params['json'] = config.post_data  #
            else:
                request_params['data'] = config.post_data  # 以表单格式发送
        # URL(url,encoded=True) 使得aiohttp其不会自动编码
        async with semaphore:
            async with getattr(session, method.lower())(**request_params) as response: #**request_params解包字典用做参数传递
                status_code = response.status
                if config.get_location and redirect:
                    redirect_url=str(response.url)
                    if redirect_url != url:
                        location=excel_text2url(redirect_url)
                try:
                    data = bytearray() # 使用 bytearray 存储读取的数据  以下方法只拿到返回包的title数据即可，而不用获取所有的返回包
                    # 通过字节流动态读取返回内容，读取到title就结束
                    # 通过二进制的正则匹配制定title的匹配规则,作者提示:要想匹配准确，正则要写好~
                    # re.IGNORECASE表示忽略大小写
                    title_pattern = re.compile(rb'<title[^>]*>(.*?)</title>', re.IGNORECASE)
                    m=0
                    title_find=False
                    # iter_chunked是aiohttp的分块儿读取的方法
                    async for chunk in response.content.iter_chunked(2048):
                        m+=1
                        data.extend(chunk)
                        # 尝试在数据块儿中查找<title>标签
                        match = title_pattern.search(data)
                        if match:
                            title_find=True
                            data = match.group(1)
                            break  # 找到标题后结束循环
                        # 不要读取完整个页面，控制读取循环4*2048数据块儿
                        if m>3:
                            if not title_find:
                                break
                    # 往后读取4次或者是页面没有2048字节的数据就提取退出。退出后校验title是否在data数据块儿中
                    if not title_find:
                        title = "未获取到标题"
                    else:
                        title=gettitle(data)
                    # 查找<title>和</title>的位置
                    # 先将二进制数据中的title包裹部分匹配拿出来在进行字符串编码，如果先将二进制数据整体转化编码之后再通过正则或者其他方法匹配很容易出现编码报错
                except Exception as e:
                    print('标题获取错误:'+ e)
                    title = '未获取到标题'
                if not str(status_code).startswith('2'):

                    print(Fore.RED + '[+]' + Fore.RESET + f"网页:{url} 状态码:" + Fore.RED + f"{status_code}" + Fore.RESET + f" 标题:{title}")
                else:
                    print(Fore.GREEN +'[+]'+ Fore.RESET+f"网页:{url} 状态码:"+Fore.GREEN+f"{status_code}"+Fore.RESET+f" 标题:{title}")
            return url, status_code,title,location
    # 如果说访问请求报错的话就执行下面的语句
    except Exception as e:
        status_code=404
        title='该页面可能无法访问！'
        # print(e)
        # 定义重试策略
        if retry < config.max_retries:  # 设置最大重试次数
            print(Fore.RED + '[-]' + Fore.RESET + f'{url} 访问失败,重试次数:{retry+1}')
            return await get_title_code(url, semaphore, method, session, redirect, retry + 1)
        else:
            print(Fore.RED + '[-]' + Fore.RESET + f'{url} 可能无法访问')
            return url, status_code, title,location


url_counts=0
# 异步生成器加载数据然后aiohttp协程任务批量请求
async def put_queue_data(data_queue, filename, method,redict):
    global url_counts
    speed=config.exploresemaphore #用户自定义的并发
    finall_semaphore=config.exploresemaphore+12 if config.exploresemaphore % 2==0 else config.exploresemaphore+11 #在此基础上可以增加一些,最好为偶数
    # 根据操作系统设置最大信号量和连接数量
    limit = config.tcp_limit
    if plat == 'win32':
        if finall_semaphore > 500:
            finall_semaphore = 500
    else:
        if finall_semaphore > 1000:
            finall_semaphore = 1000
    semaphore = asyncio.Semaphore(finall_semaphore)
    print('并发数量为'+str(speed)+'可在config.py文件当中设置')
    # 设置dns列表
    if config.dns_list:
        resolver = AsyncResolver(nameservers=config.dns_list)
    else:
        resolver = AsyncResolver()  # 默认使用系统配置的 DNS

    if config.force_close_tcp==True: #是否强制关闭tcp连接
        connector = aiohttp.TCPConnector(limit=limit,ssl=False,force_close=True,resolver=resolver,enable_cleanup_closed=True,ttl_dns_cache=60, limit_per_host=30)
    else:
        connector = aiohttp.TCPConnector(limit=limit, ssl=False, resolver=resolver, enable_cleanup_closed=True, ttl_dns_cache=60, limit_per_host=30, keepalive_timeout=20)
    #force_close=True强制关闭套接字,enable_cleanup_closed强制关闭ssl防止ssl泄露,设置ttl_dns_cache缓存为30秒
    print("超时设置:"+str(config.explore_timeout)+"秒")
    session = None  # 初始化 session
    url_count = 0 # 初始化 url_count
    try:
        timeout = aiohttp.ClientTimeout(
            total=None,  # 总超时
            sock_connect=config.explore_timeout,  # 连接超时时间
            connect=config.explore_timeout * 2 if config.force_close_tcp==False else None,
            sock_read=config.explore_timeout,  # 读取超时时间
        )  # 超时覆盖问题:https://github.com/aio-libs/aiohttp/issues/3203
        async with aiohttp.ClientSession(connector=connector,timeout=timeout) as session:
            async for urllist in read_lines(filename, speed):  # read_lines 函数提供异步生成器可以读取在读取大文件的时候减小性能消耗
                url_count += len(urllist)
                # 使用 asyncio.as_completed 来逐个处理结果
                tasks = [asyncio.create_task(get_title_code(url, semaphore, method, session, redict, 0)) for url in urllist]
                for completed in asyncio.as_completed(tasks):
                    result = await completed  # 等待单个任务完成
                    # print(result)
                    await data_queue.put(result)  # 立即将结果放入数据队列
            await data_queue.put(False) # 告知消费者任务结束
            url_counts= url_count
            print("正在保存结果文件...")
    except Exception as e:
        print("协程存活探测任务过程出错了:")
        print(e)
    finally:
        # 最后强制释放资源
        semaphore.release()
        await connector.close()  # 强制释放连接
        if session:
            await session.close()

# 主函数
async def main(nowtime,method,filename,redict):
    data_queue = asyncio.Queue(maxsize=100) #由于消息队列的初始化是在异步任务外部进行，可能会导致该消息队列被绑定到与异步任务不同的事件循环。

    print(f'正在使用{method}方法进行存活探测,请稍等一下哦~\n')
    starttime = time.time()
    process_task = asyncio.create_task(put_queue_data(data_queue, filename, method,redict))
    # await process_task
    # await data_queue.join() #数据放到队列当中，然后再写入。
    Write_task = asyncio.create_task(Write_excelandtxt(config, nowtime, method, data_queue))
    result=await asyncio.gather(process_task, Write_task) #启动消费者和生产者协程
    # counts=await Write_task
    global url_counts
    print("保存成功!")
    timecost=time.time()-starttime
    print(f'总任务共用时:{timecost:.2f}秒')
    print(f'总共存活探测url数量为{url_counts},状态码为200的url数量为{result[1]}')
    return url_counts


# 将url转化为超链接
def excel_text2url(link_url): #如果函数内部没有进行异步操作，使用 async 并不会对性能或功能产生实际影响。
    '''把一个网址字符串转换为 Excel公式，使其可以点击直接转跳'''
    if config.link_url:
        return f'=HYPERLINK("{link_url}","{link_url}")'
    else:
        return link_url


# 这里是消费者异步函数，为了确保双协议探测的情况下，在某一域名的某一协议存活的情况下，这一域名不在非存活的文件中保存，所以写入非存活文件是最后进行的，写入存活url文件是和存活探测同时进行
async def Write_excelandtxt(config, nowtime, method, data_queue):
    """写入excel和异步txt"""
    only_active_utl=set()
    else_urls_lists=[]
    activeurl = 0
    # 创建一个工作簿,设置为只写模式 占用资源小 速度更快
    wb = Workbook(write_only=True)
    # 创建两个工作表
    # 只写模式不允许直接修改 active 工作表的标题
    ws_200 = wb.create_sheet(title="200 Status")  # 创建一个标题为 "200 Status" 的工作表
    ws_else = wb.create_sheet(title="Other Status")
    pattern = r"https?://(.+)"
    try:
        # 写入标题行
        if config.get_location:
            ws_200.append(["url", "状态码", "标题","转跳到"])
        else:
            ws_200.append(["url", "状态码", "标题"])
        ws_else.append(["url", "状态码", "标题"])
        async with aiofiles.open(f'./result/{nowtime}存活探测{method}/{config.urlstxtname}.txt', 'wt', newline='',
                                 encoding='utf_8') as file_200_txt, \
                aiofiles.open(f'./result/{nowtime}存活探测{method}/{config.url_stxtname}.txt', 'wt', newline='',
                              encoding='utf_8') as file_else_txt:
            while True:
                # size=data_queue.qsize()
                # print(f"当前队列大小{size}")
                result = await data_queue.get() # 每次从队列中获取一组数据,如果没有则等待
                # 通知队列任务已完成，更新计数器，这个一般和join配合，但是这里仅为了保险起见防止bug，因为队列的计数_unfinished_tasks 属性不会随着取出而减小
                data_queue.task_done()
                # print(result)
                if result is False: #数据为false则退出循环
                    # print("退出消费者")
                    break
                # count += 1
                url, code, title,location = result  # result是个包含标题、URL和状态码的元组
                # 如果配置了状态码为200的协议唯一保存
                if config.onlysave and int(code) == 200:
                    match = re.search(pattern, url)
                    if match:
                        # 获取匹配到的内容
                        matched_url = match.group(1).strip()
                        # 检查匹配到的部分是否已经在 only_utl 列表中
                        if matched_url in only_active_utl:
                            continue
                        else:
                            # 将匹配的部分添加到 only_utl 列表中
                            only_active_utl.add(matched_url)
                # 使用 excel_text2url 函数将 URL 转换为超链接公式后写入
                url_a = excel_text2url(url)
                if int(code) == 200:
                    activeurl += 1
                    # 写入 URL 作为超链接公式
                    if config.get_location:
                        ws_200.append([url_a, code, title,location])
                    else:
                        ws_200.append([url_a, code, title])
                    await file_200_txt.write(url + "\n")
                else:
                    else_urls_lists.append([url, code, title])
                    ws_else.append([url_a, code, title])
                    await file_else_txt.write(url + "\n")
            # # 不写入存活的域名到others
            # for else_urls in else_urls_lists:
            #     domain = re.search(pattern, else_urls[0]).group(1)
            #     # 判断该子域是否在状态码为200的域名里，如果存在则不写入。
            #     if domain in only_active_utl:
            #         continue
            #     else:
            #         await file_else_txt.write(else_urls[0] + "\n")
            #         else_urls[0] = excel_text2url(else_urls[0])
            #         ws_else.append(else_urls)
            # del only_active_utl,else_urls_lists
        # 保存 Excel 文件
        wb.save(f'./result/{nowtime}存活探测{method}/{config.excelname}.xlsx')
        wb.close() #这个不需要显示 不写也行 保险起见哈哈
    except Exception as e:
        print("保存结果文件时出错:", e)
        sys.exit(0)
    return activeurl


# 方法入口
def Explore(filename,nowtime,method,redict):
    try:
        # 下面进行实际循环策略不同操作系统的设置
        if plat == 'win32':
            # 在 Windows 系统上使用 SelectorEventLoop
            policy = asyncio.WindowsSelectorEventLoopPolicy()
            asyncio.set_event_loop_policy(policy)
        else:
            # 在其他系统上使用默认的事件循环策略
            policy = asyncio.DefaultEventLoopPolicy()
            asyncio.set_event_loop_policy(policy)
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main(nowtime,method,filename,redict))
        loop.close()
    except Exception as e:
        print("协程任务启动出错!!:")
        print(e)

async def explore(filename, nowtime, method,redict):
    count = await main(nowtime, method, filename,redict)
    return count